#include "FeatureBuilder.h"

namespace LSCM
{
	FeatureBuilder::FeatureBuilder(Mesh* mesh, FeatureMetric* metric, int maxSetSize, int minFeatureLength)
		: m_mesh(mesh), m_metric(metric), m_maxStringSize(maxSetSize), m_minFeatureLength(minFeatureLength)
	{
		// 启用面法向属性，并生成面法向
		m_mesh->request_face_normals();
		m_mesh->update_face_normals();

		// 添加两个属性，分别储存特征值和状态
		m_mesh->add_property(m_featureValue);
		m_mesh->add_property(m_meshStatus);
	}

	FeatureBuilder::~FeatureBuilder()
	{
		// 完成后移除两个属性
		m_mesh->remove_property(m_featureValue);
		m_mesh->remove_property(m_meshStatus);
	}

	const std::vector<FeatureSet>& FeatureBuilder::Build()
	{
		// 获得网格和度量
		auto mesh = m_mesh;
		auto& metric = m_metric;

		// 获得边的数量
		const auto numEdges = mesh->n_edges();

		// 预留空间，并不创建元素，之后需要 push_back 填充
		std::vector<Feature> features;
		features.reserve(numEdges);

		std::cout << "Edges: " << numEdges << std::endl;

		// 遍历所有边，填充特征数组
		for (auto e_it = mesh->edges_begin(), e_end = mesh->edges_end(); e_it != e_end; e_it++)
		{
			const auto& edge = *e_it;

			// 利用二面角估计这条边的 SOD 二阶差分，它是 SOD 的补角
			auto value = metric->eval(mesh, edge);

			// 给边赋予特征值为 SOD 的补角
			m_mesh->property(m_featureValue, edge) = value;

			// 记录这个边的特征
			features.emplace_back(edge, value);

			// 将这条边对应的半边的句柄都标记为 MS_None
			for (int i = 0; i < 2; i++)
				mesh->property(m_meshStatus, mesh->halfedge_handle(*e_it, i)) = MeshStatus::MS_None;
		}

		// 根据边的特征值从大到小排序
		std::sort(features.begin(), features.end(), Feature::greater);

		auto compare = [&metric](const Feature& f)
		{
			return *metric > f.value();
		};

		// 遍历所有特征，找到第一个不大于 metric 给出的特征阈值的特征
		auto iter = std::find_if(features.begin(), features.end(), compare);

		// 修正特征数量为所有大于阈值特征的数量
		int numFeatures = 0;
		if (iter != features.end())
			numFeatures = (int)std::distance(features.begin(), iter);

		// 重置大小，相当于移除了不大于阈值的特征
		features.resize(numFeatures);

		std::cout << "features: " << numFeatures << std::endl;
		std::cout << "Threshold: " << metric->threshold() << std::endl;

		// 遍历大于阈值的特征
		int totalFeatureSize = 0;
		for (const auto& feature : features)
		{
			size_t halfLengths[] = { 0, 0 };

			// 拿到一个特征边，然后取正反两个半边进行扩张
			auto edge = feature.edge();
			FeatureSet featureSet;
			for (auto i = 0; i < 2; i++)
			{
				expand(mesh->halfedge_handle(edge, i), featureSet);

				// 记录两个半边扩张后特征集的规模
				halfLengths[i] = featureSet.size();
			}

			// 只要这一组特征数量足够多，就记录这个特征
			if (featureSet.size() > m_minFeatureLength)
			{
				// 排序特征集中的半边
				std::sort(featureSet.halfedges().begin(), featureSet.halfedges().end());

				// 把特征集中的半边及其邻域标记为 MS_FEATURE
				for (const auto& halfedge : featureSet.halfedges())
					tagHalfEdge(halfedge);

				// 记录这个特征集
				m_featureSets.push_back(featureSet);

				totalFeatureSize += featureSet.size();

				// 输出这个特征扩张后的大小，以及两个半边扩张后的规模差
				std::cout << "Feature " << feature.edge() << " : " << featureSet.size() << " = " << halfLengths[0] << " + " << halfLengths[1] - halfLengths[0] << std::endl;
			}
		}

		// 特征集的数量和特征集的平均大小
		std::cout << "Feature Sets: " << m_featureSets.size() << std::endl;
		std::cout << "Avg. Size: " << (int)(1.0 * totalFeatureSize / m_featureSets.size()) << std::endl;

		return m_featureSets;
	}

	void FeatureBuilder::expand(const Mesh::HalfedgeHandle& halfedge, FeatureSet& featureSet)
	{
		// 指定阈值，例如最多 5 条边，那么要求特征总和不小于 5 * 特征值的最小值（阈值）
		const auto threshold = m_maxStringSize * m_metric->threshold();

		// 如果特征集为空，就添加这个半边作为特征
		auto mesh = m_mesh;
		auto baseEdge = halfedge;
		if (featureSet.empty())
			featureSet.add(baseEdge, mesh->property(m_featureValue, mesh->edge_handle(baseEdge)));

		double valueTotal = 0;
		do
		{
			// 深度优先搜索
			std::vector<Mesh::HalfedgeHandle> path;
			valueTotal = depthFirst(baseEdge, path, featureSet);

			// 搜索结果为空，退出
			if (path.empty())
				break;

			// 搜索到的第 2 条边加入特征集
			baseEdge = path[1];
			featureSet.add(baseEdge, mesh->property(m_featureValue, mesh->edge_handle(baseEdge)));
		} while (valueTotal > threshold);
	}

	void FeatureBuilder::tagHalfEdge(const Mesh::HalfedgeHandle& feature)
	{
		const auto& edge = m_mesh->edge_handle(feature);
		m_mesh->property(m_meshStatus, feature) = MeshStatus::MS_Feature;

		// 获得与半边相连的前后两条半边，如果没有标记，就标记为邻域
		for (auto i = 0; i < 2; i++)
		{
			const auto& halfedge = m_mesh->halfedge_handle(edge, i);
			if (m_mesh->property(m_meshStatus, halfedge) == MeshStatus::MS_None)
				m_mesh->property(m_meshStatus, halfedge) = MeshStatus::MS_Neighbor;

			// 记录半边的起点和终点
			tagNeighborhood(m_mesh->to_vertex_handle(halfedge));
			tagNeighborhood(m_mesh->from_vertex_handle(halfedge));
		}
	}

	void FeatureBuilder::tagNeighborhood(const Mesh::VertexHandle& vertex)
	{
		// 遍历 vertex 周围的点
		for (auto v_it = m_mesh->cvv_begin(vertex), v_end = m_mesh->cvv_end(vertex); v_it != v_end; v_it++)
		{
			const auto& v = *v_it;
			for (auto e_it = m_mesh->cve_begin(v), e_end = m_mesh->cve_end(v); e_it != e_end; e_it++)
			{
				// 把前后两个半边都标记为邻域
				for (auto i = 0; i < 2; i++)
				{
					const auto& halfedge = m_mesh->halfedge_handle(*e_it, i);
					if (m_mesh->property(m_meshStatus, halfedge) == MeshStatus::MS_None)
						m_mesh->property(m_meshStatus, halfedge) = MeshStatus::MS_Neighbor;
				}
			}
		}
	}

	double FeatureBuilder::depthFirst(const Mesh::HalfedgeHandle& baseEdge, HalfedgeList& path, FeatureSet& featureSet)
	{
		auto mesh = m_mesh;

		// 通过半边句柄索引对应的状态，即储存每个半边的状态（搜索状态和网格状态）
		std::map<Mesh::HalfedgeHandle, SearchStatus> searchStatus;
		searchStatus[baseEdge] = SearchStatus(
			mesh->calc_edge_vector(baseEdge).normalized(),
			mesh->property(m_featureValue, mesh->edge_handle(baseEdge))
		);

		// 使用基础边进行初始化
		std::stack<Mesh::HalfedgeHandle> dfs;
		dfs.push(baseEdge);

		auto maxEdge = Mesh::HalfedgeHandle();
		auto maxValue = 0.0f;

		// 对栈中的元素进行深度搜索
		while (!dfs.empty())
		{
			// 弹出栈顶半边作为起点
			auto parentEdge = dfs.top();
			dfs.pop();

			// 跳过已经访问的半边
			if (mesh->property(m_meshStatus, parentEdge) != MeshStatus::MS_None)
				continue;

			// 检查这个半边是否在分支中，如果在就跳过
			auto& parentStatus = searchStatus[parentEdge];
			if (parentStatus.included)
				continue;

			// 现在加入分支
			parentStatus.included = true;

			// 如果深度达到要求，跳过
			if (parentStatus.depth >= m_maxStringSize)
				continue;

			// 获得半边终点，遍历以它为起点的半边 oh - outgoing halfedge
			const auto& v = mesh->to_vertex_handle(parentEdge);
			for (auto e_it = mesh->voh_begin(v), e_end = mesh->voh_end(v); e_it != e_end; e_it++)
			{
				// 这个半边作为子边
				auto childEdge = *e_it;

				// 跳过已经标记的边还有已经在当前特征集中的边
				if (mesh->property(m_meshStatus, childEdge) != MeshStatus::MS_None || featureSet.contains(childEdge))
					continue;

				// 检查这个半边是否已经搜索过，如果搜索过就使用之前的数据，否则新建数据
				// 检查是否已经包含在分支中，如果是就跳过
				auto childStatus = searchStatus.count(childEdge) > 0 ? searchStatus[childEdge] : SearchStatus();
				if (childStatus.included)
					continue;

				// 获得该半边的方向，是否前面半边同向（夹角小于 90）
				bool oppsite = false;
				auto edge = parentEdge;
				auto dir = mesh->calc_edge_vector(childEdge).normalized();
				while (edge != Mesh::HalfedgeHandle())
				{
					const auto& status = searchStatus[edge];
					auto angle = OpenMesh::rad_to_deg(std::acos(OpenMesh::dot(dir, status.dir)));

					if (angle >= 90.0f)
					{
						oppsite = true;
						break;
					}

					edge = status.parent;
				}

				// 如果不同向，就跳过
				if (oppsite)
					continue;

				// 现在可以记录该半边
				dfs.push(childEdge);

				// 更新状态：1.如果没有记录，这就是第一次记录；2.如果记录过，由于深度优先，它被更新到与一个更深的半边相连
				childStatus.parent = parentEdge;
				childStatus.dir = dir;
				childStatus.valueTotal = parentStatus.valueTotal + mesh->property(m_featureValue, mesh->edge_handle(childEdge));
				childStatus.depth = parentStatus.depth + 1;

				searchStatus[childEdge] = childStatus;

				// 刷新特征和的记录，主要用来更新最大边
				if (childStatus.valueTotal > maxValue)
				{
					maxEdge = childEdge;
					maxValue = childStatus.valueTotal;
				}
			}
		}

		// 从最大边开始，向上延伸获得返回路径
		auto edge = maxEdge;
		while (edge != Mesh::HalfedgeHandle())
		{
			path.push_back(edge);
			const auto& status = searchStatus[edge];
			edge = status.parent;
		}
		std::reverse(path.begin(), path.end());

		return maxValue;
	}
}